;;; An experiment framework

(in-package :experiment)

(defun current-date ()
  "Return the current date and time as a formatted string."
  (multiple-value-bind (s min h d m y) (decode-universal-time (get-universal-time))
    (format nil "~2,'0d:~2,'0d:~2,'0d ~2,'0d/~2,'0d/~d" h min s d m y)))

(defgeneric dump-screen (experiment state))
(defgeneric handle-event (experiment event))

(defvar *default-checkpointp* nil
  "The default value of a state's checkpointp slot.")

(defvar *experiment* nil
  "dynamically bound to the current experiment instance.")

(defclass experiment (gwee:switchable-layout)
  ((id :initarg :id :accessor experiment-id)
   (session-name :initarg :session-name :accessor experiment-session-name)
   (reward :initform 0 :initarg :reward :accessor reward)
   (log-file :initarg :log-file :accessor experiment-log-file)
   (log-stream :initform nil :accessor log-stream)
   (event-file :initarg :event-file :accessor experiment-event-file)
   (event-stream :initform nil :accessor event-stream)
   (suspend-file :initarg :suspend-file :accessor experiment-suspend-file)
   (state-stack :initform nil :accessor state-stack)
   (start-time :accessor experiment-start-time)
   (disable-timers-p :initform nil :accessor experiment-disable-timers-p :documentation
"Models don't want timers interfering with their state. So set this to
T and let the model advance the timers manually.")
;;    (state-stack :initform nil :accessor experiment-state-stack)
;;    (state-tree :initform nil :accessor experiment-state-tree)
   (global-data :initform (make-hash-table))
   (options :initarg :options :accessor options)
   (fullscreen-hook :initform nil :initarg :fullscreen-hook :accessor fullscreen-hook)
   (time-function :initarg nil :initarg :time-function :accessor time-function)))

(defmethod shared-initialize :after ((exp experiment) slots &key global-data)
  (declare (ignore slots))
  (setf (slot-value exp 'global-data) (u:plist-to-hash global-data)))

(defgeneric current-state (experiment))
(defgeneric parent-state (experiment))
(defgeneric restore-experiment (experiment))
(defgeneric save-experiment (experiment))
(defgeneric dump-experiment-slots (experiment)
  (:method-combination append :most-specific-last))
(defgeneric global-data (experiment key))
(defgeneric (setf global-data) (value experiment key))

(defmethod current-state ((exp experiment))
  (first (state-stack exp)))

(defmethod parent-state ((exp experiment))
  (second (state-stack exp)))

(defmethod dump-experiment-slots append ((exp experiment))
  `(:reward ,(reward exp)
    :session-name ,(experiment-session-name exp)
    :global-data ,(u:hash-to-plist (slot-value exp 'global-data))))

(defmethod save-experiment ((exp experiment))
  (unless *skip-log-file*
    (ensure-directories-exist (experiment-suspend-file exp))
    (with-open-file (s (experiment-suspend-file exp)
                       :direction :output
                       :if-exists :supersede
                       :if-does-not-exist :create
                       :external-format :utf-8)
      (let ((*print-right-margin* 120))
        (print `(,(type-of exp)
                  ,@(dump-experiment-slots exp)
                  :states ,(state-dump exp (car (last (state-stack exp)))))
               s)))))

(defmethod restore-experiment ((exp experiment))
  "This method has side-effects!"
  (with-open-file (s (experiment-suspend-file exp)
                     :external-format :utf-8)
    (let ((e (read s)))
      (assert (and (listp e)
                   (eq (first e) 'experiment)))
      (prog1
          (getf (cdr e) :states)
        (remf (cdr e) :states)
        ;; side effects!
        (apply 'reinitialize-instance exp (cdr e))))))

(defmethod global-data ((exp experiment) key)
  (gethash key (slot-value exp 'global-data)))

(defmethod (setf global-data) (value (exp experiment) key)
  (setf (gethash key (slot-value exp 'global-data)) value))

(defmethod get-option-value ((exp experiment) key)
  (let ((o (find key (options exp) :key 'option-name)))
    (when o
      (values (exp:option-value o) t))))

(defmacro with-current-state ((exp state) &body body)
  `(progn
     (push ,state (state-stack ,exp))
     (unwind-protect
          (progn ,@body)
       (assert (eq (first (state-stack ,exp)) ,state))
       (pop (state-stack ,exp)))))

(defmethod handle-event (experiment event)
  (state-dispatch-event experiment (current-state experiment) event))

(defmacro dispatch-event-fn (experiment state event &rest args)
  (let ((junk (gensym "JUNK")))
    `(lambda (&rest ,junk)
       (declare (ignore ,junk))
       (state-dispatch-event ,experiment ,state
                             (make-instance ',event ,@args)))))

(defun wire-button (button source &key (exp *experiment*) (state (current-state exp)))
  (setf (gwee:click-callback button)
        (dispatch-event-fn exp state exp:click-event :source source)))

(defmethod gwee:key-press-event ((view experiment) key)
  (if (eq (gwee::key-sym key) :sdl-key-escape)
      (let ((mode (gwee:toggle-fullscreen)))
        (when (fullscreen-hook *experiment*)
          (funcall (fullscreen-hook *experiment*) mode)))
      (state-dispatch-event view (current-state view)
			    (make-instance 'key-press-event :key key))))

(defmethod gwee:mouse-press-event ((view experiment) button x y)
  (or (call-next-method)
      (progn
        (state-dispatch-event view (current-state view)
                              (make-instance 'click-event :source (gwee:find-subview-at view x y) :x x :y y))
        t)))

(defun compile-experiment (experiment states &key (default-checkpointp *default-checkpointp*))
  "If states is a function then call it with EXPERIMENT as an
argument. Sometimes generating the states depends on the subject ID."
  (let ((*default-checkpointp* default-checkpointp))
    (when (functionp states)
      (setf states (funcall states experiment)))
    (if (> (length states) 1)
        (compile-state experiment 'subsequence states)
        (compile-state experiment (car (car states)) (cdr (car states))))))

(defun experiment-log-start (exp extra &key random-seed session (version "Unknown") options)
  (setf (experiment-start-time exp) (funcall (time-function exp)))
  (log-event-start version session (experiment-id exp) options random-seed)
  (apply 'log exp 'start
         :id (experiment-id exp)
         :time (experiment-start-time exp)
         :date (current-date)
         :session session
         :random-seed random-seed
         :version version
         (append (mapcan (lambda (o) (list (option-name o) (option-value o))) options)
                 extra)))

(defun experiment-close-streams (exp)
  (close-log-stream exp)
  (close-event-stream exp))

(defun experiment-log-end (exp &key session (version "Unknown"))
  (log exp 'end
       :id (experiment-id exp)
       :session session
       :version version)
  (experiment-close-streams exp))

(defstruct main-screen
  sessions options)

(defvar *main-screen* nil
  "Dynamically bound to the list of sessions and options available when in the main screen.")

(defun experiment-main-loop (&key options id log-template event-template default-checkpointp session sessions version fullscreen-hook time-function)
  (let* ((*experiment* (make-instance 'experiment
                                      :id id
                                      :session-name session
                                      :log-file (get-log-file log-template id session)
                                      :event-file (get-event-file event-template id session)
                                      :suspend-file (get-suspend-file id)
                                      :parent gwee:*top-view*
                                      :x 0 :y 0
                                      :width :maximum :height :maximum
                                      :options options
                                      :fullscreen-hook fullscreen-hook
                                      :time-function time-function))
         (*default-checkpointp* default-checkpointp)
         (seed (+ (sxhash id) (sxhash session))))
    (randomize seed)
    (let ((states (cond
                    ((eq session :resume)
                     (prog1
                         (compile-experiment *experiment* (list (restore-experiment *experiment*)))
                       ;; We need to fix the log and event files in case the session name is used in their file names
                       (setf (experiment-log-file *experiment*) (get-log-file log-template (experiment-id *experiment*) (experiment-session-name *experiment*))
                             (experiment-event-file *experiment*) (get-event-file log-template (experiment-id *experiment*) (experiment-session-name *experiment*)))))
                    ((u:funcallable (second (assoc session sessions :test 'equal)))
                     (record-warnings
                      (funcall (second (assoc session sessions :test 'equal)) *experiment*)))
                    (t
                     (compile-experiment *experiment* (cdr (assoc session sessions)))))))
      (unwind-protect
           (progn
             (experiment-log-start *experiment* nil
                                   :random-seed seed
                                   :options options
                                   :session session
                                   :version version)
             (state-run *experiment* states)
             (experiment-log-end *experiment*
                                 :session session
                                 :version version))
        (experiment-close-streams *experiment*)))))

(defun end-menu (no-video-dump)
  (let ((menu (make-instance 'gwee:column-layout
                             :background (gwee:gray 220)
                             :parent gwee:*top-view*
                             :padding 10
                             :width 200
                             :height :minimum
                             :x :center
                             :y :bottom
                             :subviews (delete nil
                                               (list
                                                (make-instance 'gwee:text-area :data "Session is over"
                                                                               :foreground (gwee:gray 0)
                                                                               :background nil)
                                                (make-instance 'gwee:button :data "Main Menu"
                                                                            :click-callback (lambda (b)
                                                                                              (declare (ignore b))
                                                                                              (throw 'choose :restart)))
                                                (unless no-video-dump
                                                  (make-instance 'gwee:button :data "Video Dump"
                                                                              :disabled-p (find :powerpc *features*)
                                                                              :click-callback (lambda (b)
                                                                                                (declare (ignore b))
                                                                                                (throw 'choose :video)))))))))
    (setf (gwee:view-needs-update-p gwee::*top-view*) t)
    (unwind-protect
         (catch 'choose
           (gwee:dispatch-events))
      (gwee:delete-view menu))))

(defun run-experiment (sessions &rest keys &key resizable (title "Experiment") (version "Unknown" version-supplied) options startup-message (login (make-instance 'login)) (width 1024) (height 768) (background (gwee:gray 255)) (foreground (gwee:gray 0)) (log-template "log-~a.txt") (event-template "events-~a.txt") log model-server start-callback (default-checkpointp t) of-interest skip-log-file disable-resume disable-video-dump (default-id "") (hide-cursor nil) (fullscreen-hook nil) (time-function #'get-internal-real-time))
  (unless version-supplied
    (warn "!!! No version was supplied !!!"))
  (gwee:with-gui (width height
                        :title title
                        :font "sans.ttf"
                        :background background
                        :foreground foreground
                        :resizable resizable
                        :event-hook 'log-event-hook
                        :timer-hook 'log-timer-hook)
    (if skip-log-file
        (setf *skip-log-file* t)
        (when (read-only-fs-p)
          (gwee:display-message "You are running from a read-only file system so a log file will not be written."
                                :width 500)
          (sdl:fill-surface (gwee:gray 0))
          (gwee::update-display)))
    (when startup-message
      (gwee:display-message startup-message :width 500)
      (sdl:fill-surface (gwee:gray 0))
      (gwee::update-display))
    (when model-server
      (let ((ret (model-server:start)))
        (when (typep ret 'condition)
          (gwee:display-message (format nil "Failed to start model server: ~a" ret))
          (sdl:fill-surface (gwee:gray 0))
          (gwee::update-display))))
    ;; It's handy in sandbox modes to be able to jump back and do it
    ;; all over again.
    (loop
      (catch 'restart
        (gwee:cancel-all-timers)
        (setf (gwee:view-subviews gwee:*top-view*) nil)
        (multiple-value-bind (id session) (let ((*main-screen* (make-main-screen :sessions sessions :options options)))
                                            (get-subject-from-user login default-id sessions options disable-resume disable-video-dump event-template))
          (if (eq session :videos)
              (progn
                (if (probe-file (get-event-file id event-template))
                    ;; (handler-case
                    (replay-events-from-file id event-template sessions nil
                                             :of-interest of-interest 
                                             :dry-run nil)
                    ;; (error (c)
                    ;;   (gwee:display-message (format nil "An error occured during replay: ~a" c))))
                    (gwee:display-message (format nil "Cannot find event file for subject ~d" id)))
                (throw 'restart nil))
              (progn
                (sdl:show-cursor (not hide-cursor))
                (experiment-main-loop :session session
                                      :id id
                                      :options options
                                      :log-template log-template
                                      :event-template event-template
                                      :default-checkpointp default-checkpointp
                                      :sessions sessions
                                      :version version
                                      :fullscreen-hook fullscreen-hook
                                      :time-function time-function)
                (sdl:show-cursor t)
                (when (gwee::fullscreen-p)
                  (gwee:toggle-fullscreen))
                (ecase (end-menu disable-video-dump)
                  (:restart
                   (throw 'restart nil))
                  (:video
                   (replay-events-from-file id event-template sessions t
                                            :of-interest of-interest
                                            :dry-run nil))))))))))
